#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import pymysql.cursors
import csv
import pymongo
import requests
from tqdm import tqdm
import configparser
import click
import datetime
import os


@click.group()
def main():
    """
    NCBI Virus database mirroring toolkits, includes metadata/fasta downloading and database updating.
    """
    pass


def connect_mysql(configpath: str):
    config = configparser.RawConfigParser()
    config.read(configpath)

    host = config.get("mysql", "host")
    port = int(config.get("mysql", "port"))
    user = config.get("mysql", "user")
    password = config.get("mysql", "password")
    db = config.get("mysql", "db")
    connection = pymysql.connect(host=host,
                                 port=port,
                                 user=user,
                                 password=password,
                                 db=db,
                                 charset='utf8',
                                 cursorclass=pymysql.cursors.DictCursor)
    return connection


def connect_mongodb(configpath: str):
    config = configparser.RawConfigParser()
    config.read(configpath)

    mongo_host = config.get("mongodb", "host")
    mongo_port = int(config.get("mongodb", "port"))
    mongo_db = config.get("mongodb", "db")
    mongo_collection = config.get("mongodb", "collection")
    client = pymongo.MongoClient(mongo_host, mongo_port)
    db = client[mongo_db]
    collection = db[mongo_collection]
    return client, collection


def download_file(url, name):
    resp = requests.get(url=url, stream=True)
    # stream=True的作用是仅让响应头被下载，连接保持打开状态，
    # content_size = int(resp.headers['Content-Length']) / 1024  # 确定整个安装包的大小
    with open(name, "wb") as f:
        # print("Total file size: ", content_size, 'k，start downloading...')
        print("Start downloading " + name + " ...")
        for data in tqdm(iterable=resp.iter_content(1024), unit="k", desc=name):
            # 调用iter_content，一块一块的遍历要下载的内容，搭配stream=True，此时才开始真正的下载
            # iterable：可迭代的进度条 total：总的迭代次数 desc：进度条的前缀
            f.write(data)
        print(name + " is downloaded.")


@main.command("download")
@click.option("--source", "-s", help="refseq or genbank", required=True,
              type=click.Choice(['refseq', 'genbank'], case_sensitive=False), default="refseq", show_default=True)
@click.option("--type", "-t", help="nucleotide or protein", required=True,
              type=click.Choice(['nucleotide', 'protein'], case_sensitive=False), default="nucleotide",
              show_default=True)
@click.option("--data", "-d", help="fasta, metadata, accession, all", required=True,
              type=click.Choice(['fasta', 'metadata', 'accession', 'all'], case_sensitive=False), default="all",
              show_default=True)
@click.argument("config_file", type=click.Path(exists=True))
def download_files(**kwargs):
    """
        NCBI Virus database download toolkits.
    """
    source = kwargs["source"]
    data_type = kwargs["type"]
    data = kwargs["data"]

    now_date = datetime.datetime.now().strftime("%Y%m%d")
    dirs = "ncbi_virus_" + now_date
    if not os.path.exists(dirs):
        os.makedirs(dirs)

    configpath = kwargs["config_file"]
    config = configparser.RawConfigParser()
    config.read(configpath)
    if source == "refseq":
        if data_type == "nucleotide":
            if data == "all":
                download_file(config.get("urls", "ALL_VIRUS_REFSEQ_NUC_FASTA"),
                              os.path.join(dirs, "ncbi_virus_refseq_nucleotide.fasta"))
                download_file(config.get("urls", "ALL_VIRUS_REFSEQ_NUC_CSV"),
                              os.path.join(dirs, "ncbi_virus_refseq_nucleotide.csv"))
                download_file(config.get("urls", "DNA_VIRUS_REFSEQ_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_nucleotide_dna.acc"))
                download_file(config.get("urls", "RNA_VIRUS_REFSEQ_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_nucleotide_rna.acc"))
                download_file(config.get("urls", "DNA_RNA_VIRUS_REFSEQ_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_nucleotide_dnarna.acc"))
                download_file(config.get("urls", "UNCLASSIFIED_VIRUS_REFSEQ_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_nucleotide_unclassified.acc"))
            elif data == "fasta":
                download_file(config.get("urls", "ALL_VIRUS_REFSEQ_NUC_FASTA"),
                              os.path.join(dirs, "ncbi_virus_refseq_nucleotide.fasta"))
            if data == "metadata":
                download_file(config.get("urls", "ALL_VIRUS_REFSEQ_NUC_CSV"),
                              os.path.join(dirs, "ncbi_virus_refseq_nucleotide.csv"))
            if data == "accession":
                download_file(config.get("urls", "DNA_VIRUS_REFSEQ_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_nucleotide_dna.acc"))
                download_file(config.get("urls", "RNA_VIRUS_REFSEQ_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_nucleotide_rna.acc"))
                download_file(config.get("urls", "DNA_RNA_VIRUS_REFSEQ_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_nucleotide_dnarna.acc"))
                download_file(config.get("urls", "UNCLASSIFIED_VIRUS_REFSEQ_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_nucleotide_unclassified.acc"))
        elif data_type == "protein":
            if data == "all":
                download_file(config.get("urls", "ALL_VIRUS_REFSEQ_PRO_FASTA"),
                              os.path.join(dirs, "ncbi_virus_refseq_protein.fasta"))
                download_file(config.get("urls", "ALL_VIRUS_REFSEQ_PRO_CSV"),
                              os.path.join(dirs, "ncbi_virus_refseq_protein.csv"))
                download_file(config.get("urls", "DNA_VIRUS_REFSEQ_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_protein_dna.acc"))
                download_file(config.get("urls", "RNA_VIRUS_REFSEQ_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_protein_rna.acc"))
                download_file(config.get("urls", "DNA_RNA_VIRUS_REFSEQ_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_protein_dnarna.acc"))
                download_file(config.get("urls", "UNCLASSIFIED_VIRUS_REFSEQ_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_protein_unclassified.acc"))
            elif data == "fasta":
                download_file(config.get("urls", "ALL_VIRUS_REFSEQ_PRO_FASTA"),
                              os.path.join(dirs, "ncbi_virus_refseq_protein.fasta"))
            if data == "metadata":
                download_file(config.get("urls", "ALL_VIRUS_REFSEQ_PRO_CSV"),
                              os.path.join(dirs, "ncbi_virus_refseq_protein.csv"))
            if data == "accession":
                download_file(config.get("urls", "DNA_VIRUS_REFSEQ_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_protein_dna.acc"))
                download_file(config.get("urls", "RNA_VIRUS_REFSEQ_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_protein_rna.acc"))
                download_file(config.get("urls", "DNA_RNA_VIRUS_REFSEQ_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_protein_dnarna.acc"))
                download_file(config.get("urls", "UNCLASSIFIED_VIRUS_REFSEQ_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_refseq_protein_unclassified.acc"))
    elif source == "genbank":
        if data_type == "nucleotide":
            if data == "all":
                download_file(config.get("urls", "ALL_VIRUS_GENBANK_NUC_FASTA"),
                              os.path.join(dirs, "ncbi_virus_genbank_nucleotide.fasta"))
                download_file(config.get("urls", "ALL_VIRUS_GENBANK_NUC_CSV"),
                              os.path.join(dirs, "ncbi_virus_genbank_nucleotide.csv"))
                download_file(config.get("urls", "DNA_VIRUS_GENBANK_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_nucleotide_dna.acc"))
                download_file(config.get("urls", "RNA_VIRUS_GENBANK_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_nucleotide_rna.acc"))
                download_file(config.get("urls", "DNA_RNA_VIRUS_GENBANK_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_nucleotide_dnarna.acc"))
                download_file(config.get("urls", "UNCLASSIFIED_VIRUS_GENBANK_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_nucleotide_unclassified.acc"))
            elif data == "fasta":
                download_file(config.get("urls", "ALL_VIRUS_GENBANK_NUC_FASTA"),
                              os.path.join(dirs, "ncbi_virus_genbank_nucleotide.fasta"))
            if data == "metadata":
                download_file(config.get("urls", "ALL_VIRUS_GENBANK_NUC_CSV"),
                              os.path.join(dirs, "ncbi_virus_genbank_nucleotide.csv"))
            if data == "accession":
                download_file(config.get("urls", "DNA_VIRUS_GENBANK_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_nucleotide_dna.acc"))
                download_file(config.get("urls", "RNA_VIRUS_GENBANK_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_nucleotide_rna.acc"))
                download_file(config.get("urls", "DNA_RNA_VIRUS_GENBANK_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_nucleotide_dnarna.acc"))
                download_file(config.get("urls", "UNCLASSIFIED_VIRUS_GENBANK_NUC_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_nucleotide_unclassified.acc"))
        elif data_type == "protein":
            if data == "all":
                download_file(config.get("urls", "ALL_VIRUS_GENBANK_PRO_FASTA"),
                              os.path.join(dirs, "ncbi_virus_genbank_protein.fasta"))
                download_file(config.get("urls", "ALL_VIRUS_GENBANK_PRO_CSV"),
                              os.path.join(dirs, "ncbi_virus_genbank_protein.csv"))
                download_file(config.get("urls", "DNA_VIRUS_GENBANK_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_protein_dna.acc"))
                download_file(config.get("urls", "RNA_VIRUS_GENBANK_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_protein_rna.acc"))
                download_file(config.get("urls", "DNA_RNA_VIRUS_GENBANK_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_protein_dnarna.acc"))
                download_file(config.get("urls", "UNCLASSIFIED_VIRUS_GENBANK_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_protein_unclassified.acc"))
            elif data == "fasta":
                download_file(config.get("urls", "ALL_VIRUS_GENBANK_PRO_FASTA"),
                              os.path.join(dirs, "ncbi_virus_genbank_protein.fasta"))
            if data == "metadata":
                download_file(config.get("urls", "ALL_VIRUS_GENBANK_PRO_CSV"),
                              os.path.join(dirs, "ncbi_virus_genbank_protein.csv"))
            if data == "accession":
                download_file(config.get("urls", "DNA_VIRUS_GENBANK_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_protein_dna.acc"))
                download_file(config.get("urls", "RNA_VIRUS_GENBANK_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_protein_rna.acc"))
                download_file(config.get("urls", "DNA_RNA_VIRUS_GENBANK_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_protein_dnarna.acc"))
                download_file(config.get("urls", "UNCLASSIFIED_VIRUS_GENBANK_PRO_ACC"),
                              os.path.join(dirs, "ncbi_virus_genbank_protein_unclassified.acc"))


@main.command("analyze")
@click.option("--source", "-s", help="refseq or genbank", required=True,
              type=click.Choice(['refseq', 'genbank'], case_sensitive=False), default="refseq", show_default=True)
@click.option("--type", "-t", help="nucleotide or protein", required=True,
              type=click.Choice(['nucleotide', 'protein'], case_sensitive=False), default="nucleotide",
              show_default=True)
@click.argument("config_file", type=click.Path(exists=True))
@click.argument("data_path", type=click.Path(exists=True))
def update_database(**kwargs):
    """
        Analyze download files and update local MySQL database.
    """
    source = kwargs["source"]
    data_type = kwargs["type"]
    configpath = kwargs["config_file"]
    datapath = kwargs["data_path"]

    if source == "refseq":
        if data_type == "nucleotide":
            table = "virus_refseq_nucleotide"
            update_ncbi_virus(configpath, os.path.join(datapath, "ncbi_virus_refseq_nucleotide.csv"), table)
            fill_virus_lineage(configpath, table)
            fill_virus_class(configpath, datapath, table)
        elif data_type == "protein":
            table = "virus_refseq_protein"
            update_ncbi_virus(configpath, os.path.join(datapath, "ncbi_virus_refseq_protein.csv"), table)
            fill_virus_lineage(configpath, table)
            fill_virus_class(configpath, datapath, table)
    elif source == "genbank":
        if data_type == "nucleotide":
            table = "virus_genbank_nucleotide"
            update_ncbi_virus(configpath, os.path.join(datapath, "ncbi_virus_genbank_nucleotide.csv"), table)
            fill_virus_lineage(configpath, table)
            fill_virus_class(configpath, datapath, table)
        elif data_type == "protein":
            table = "virus_genbank_protein"
            update_ncbi_virus(configpath, os.path.join(datapath, "ncbi_virus_genbank_protein.csv"), table)
            fill_virus_lineage(configpath, table)
            fill_virus_class(configpath, datapath, table)


def update_ncbi_virus(configpath: str, filepath: str, table: str):
    print("Start analyzing %s and update table %s ..." % (filepath, table))
    connection = connect_mysql(configpath)

    try:
        cursor = connection.cursor()
        init_sql = "TRUNCATE `" + table + "`;"
        cursor.execute(init_sql)
        connection.commit()
        sql = ""
        if table.endswith("nucleotide"):
            sql = "INSERT INTO `" + table + "` VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) ON DUPLICATE KEY UPDATE details=%s, release_date=%s,species=%s, genus=%s, family=%s,class=%s,lineage=%s,length=%s,source=%s,completeness=%s,genotype=%s,genome_region=%s,segment=%s,country=%s,host=%s,isolation_source=%s,collection_date=%s,biosample=%s,genbank_title=%s"
        elif table.endswith("protein"):
            sql = "INSERT INTO `" + table + "` VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s) ON DUPLICATE KEY UPDATE details=%s, release_date=%s,species=%s, genus=%s, family=%s,class=%s,lineage=%s,length=%s,genotype=%s,genome_region=%s,segment=%s,protein=%s,country=%s,host=%s,isolation_source=%s,collection_date=%s,biosample=%s,genbank_title=%s"
        infile = open(filepath, "r")
        infile_csv = csv.reader(infile)
        headers = next(infile_csv)
        for row in tqdm(infile_csv, desc="Virus"):
            items = dict(zip(headers, row))
            data = gen_data_update(items, table)
            cursor.execute(sql, data)
            connection.commit()
        cursor.close()
    finally:
        connection.close()


def gen_data_update(row: dict, table: str):
    data1 = []
    data2 = []
    data1.append(None)
    if row["Accession"] != "":
        data1.append(row["Accession"])
    else:
        data1.append(None)
    if row["Details"] != "":
        data1.append(row["Details"])
        data2.append(row["Details"])
    else:
        data1.append(None)
        data2.append(None)
    if row["Release_Date"] != "":
        release_date = row["Release_Date"].split("T")[0]
        data1.append(release_date)
        data2.append(release_date)
    else:
        data1.append(None)
        data2.append(None)
    if row["Species"] != "":
        data1.append(row["Species"])
        data2.append(row["Species"])
    else:
        data1.append(None)
        data2.append(None)
    if row["Genus"] != "":
        data1.append(row["Genus"])
        data2.append(row["Genus"])
    else:
        data1.append(None)
        data2.append(None)
    if row["Family"] != "":
        data1.append(row["Family"])
        data2.append(row["Family"])
    else:
        data1.append(None)
        data2.append(None)
    data1.append(None)
    data2.append(None)
    data1.append(None)
    data2.append(None)
    if row["Length"] != "":
        length = int(row["Length"])
        data1.append(length)
        data2.append(length)
    else:
        data1.append(None)
        data2.append(None)
    if table.endswith("nucleotide"):
        if row["Nuc._Completeness"] != "":
            nuc_completeness = row["Nuc._Completeness"].split(", ")
            if len(nuc_completeness) == 2:
                data1.append(nuc_completeness[0])
                data1.append(nuc_completeness[1])
                data2.append(nuc_completeness[0])
                data2.append(nuc_completeness[1])
            else:
                data1.append(nuc_completeness[0])
                data2.append(nuc_completeness[0])
                data1.append(None)
                data2.append(None)
        else:
            data1.append(None)
            data2.append(None)
            data1.append(None)
            data2.append(None)
    if row["Genotype"] != "":
        data1.append(row["Genotype"])
        data2.append(row["Genotype"])
    else:
        data1.append(None)
        data2.append(None)
    if row["Genome_Region"] != "":
        data1.append(row["Genome_Region"])
        data2.append(row["Genome_Region"])
    else:
        data1.append(None)
        data2.append(None)
    if row["Segment"] != "":
        data1.append(row["Segment"])
        data2.append(row["Segment"])
    else:
        data1.append(None)
        data2.append(None)
    if table.endswith("protein"):
        if row["Protein"] != "":
            data1.append(row["Protein"])
            data2.append(row["Protein"])
        else:
            data1.append(None)
            data2.append(None)
    if row["Country"] != "":
        data1.append(row["Country"])
        data2.append(row["Country"])
    else:
        data1.append(None)
        data2.append(None)
    if row["Host"] != "":
        data1.append(row["Host"])
        data2.append(row["Host"])
    else:
        data1.append(None)
        data2.append(None)
    if row["Isolation_Source"] != "":
        data1.append(row["Isolation_Source"])
        data2.append(row["Isolation_Source"])
    else:
        data1.append(None)
        data2.append(None)
    if row["Collection_Date"] != "":
        collection_date = row["Collection_Date"]
        if len(collection_date) == 4:
            collection_date = collection_date + "-01-01"
        elif len(collection_date) == 7:
            collection_date = collection_date + "-01"
        data1.append(collection_date)
        data2.append(collection_date)
    else:
        data1.append(None)
        data2.append(None)
    if row["BioSample"] != "":
        data1.append(row["BioSample"])
        data2.append(row["BioSample"])
    else:
        data1.append(None)
        data2.append(None)
    if row["GenBank_Title"] != "":
        data1.append(row["GenBank_Title"])
        data2.append(row["GenBank_Title"])
    else:
        data1.append(None)
        data2.append(None)
    return data1 + data2


def fill_virus_lineage(configpath: str, table: str):
    print("Start extracting lineage from NCBI taxonomy and update table %s ..." % table)
    connection = connect_mysql(configpath)
    client, collection = connect_mongodb(configpath)

    try:
        cursor = connection.cursor()
        sql = "SELECT `virus_id`, `species` FROM `" + table + "` WHERE `species` IS NOT NULL"
        cursor.execute(sql)
        results = cursor.fetchall()
        for result in tqdm(results, desc="Virus"):
            species = result["species"]
            lineage = recursive_head_lineage(species, collection)
            if lineage is not None:
                sql2 = "UPDATE `" + table + "` SET lineage=%s WHERE `virus_id`=%s"
                cursor.execute(sql2, (lineage, result["virus_id"]))
                connection.commit()
        cursor.close()

    finally:
        connection.close()
        client.close()


def recursive_head_lineage(name: str, collection: pymongo.collection.Collection):
    lineage = [name]
    current_name = name
    current_id = get_taxon_id(name, collection)
    if current_id is None:
        return None
    while current_id != "10239":
        current_id, current_name = parent_taxon(current_id, collection)
        if current_id is None:
            return None
        lineage.append(current_name)
    lineage = lineage[-2::-1]
    lineage = "; ".join(lineage)
    return lineage


def get_taxon_id(name: str, collection: pymongo.collection.Collection):
    result = collection.find_one({"dms_name": name}, {"dms_id": 1})
    if result is None:
        return None
    else:
        return result["dms_id"]


def parent_taxon(taxid: str, collection: pymongo.collection.Collection):
    result = collection.find_one({"dms_id": taxid}, {"parent_tax_id": 1})
    if result is None:
        return None, None
    else:
        result2 = collection.find_one({"dms_id": result["parent_tax_id"]["id"]}, {"dms_name": 1})
        return result["parent_tax_id"]["id"], result2["dms_name"]


def fill_virus_class(configpath: str, datapath: str, table: str):
    connection = connect_mysql(configpath)

    try:
        cursor = connection.cursor()
        if table == "virus_refseq_nucleotide":
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_refseq_nucleotide_dna.acc"), table,
                         "DNA viruses")
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_refseq_nucleotide_rna.acc"), table,
                         "RNA viruses")
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_refseq_nucleotide_dnarna.acc"), table,
                         "DNA/RNA viruses")
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_refseq_nucleotide_unclassified.acc"),
                         table, "Unclassified viruses")
        elif table == "virus_refseq_protein":
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_refseq_protein_dna.acc"), table,
                         "DNA viruses")
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_refseq_protein_rna.acc"), table,
                         "RNA viruses")
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_refseq_protein_dnarna.acc"), table,
                         "DNA/RNA viruses")
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_refseq_protein_unclassified.acc"),
                         table, "Unclassified viruses")
        elif table == "virus_genbank_nucleotide":
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_genbank_nucleotide_dna.acc"), table,
                         "DNA viruses")
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_genbank_nucleotide_rna.acc"), table,
                         "RNA viruses")
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_genbank_nucleotide_dnarna.acc"), table,
                         "DNA/RNA viruses")
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_genbank_nucleotide_unclassified.acc"),
                         table, "Unclassified viruses")
        elif table == "virus_genbank_protein":
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_genbank_protein_dna.acc"), table,
                         "DNA viruses")
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_genbank_protein_rna.acc"), table,
                         "RNA viruses")
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_genbank_protein_dnarna.acc"), table,
                         "DNA/RNA viruses")
            update_class(connection, cursor, os.path.join(datapath, "ncbi_virus_genbank_protein_unclassified.acc"),
                         table, "Unclassified viruses")
        cursor.close()
    finally:
        connection.close()


def update_class(connection: pymysql.Connection, cursor: pymysql.cursors.Cursor, filepath: str, table: str,
                 virus_class: str):
    print("Start analyzing %s and update virus class ..." % filepath)
    count = 0
    accession = []
    infile = open(filepath, "r")
    for line in tqdm(infile, desc="Virus"):
        line = line.rstrip()
        accession.append(line)
        count += 1
        if count == 500:
            accession_str = "','".join(accession)
            sql = "UPDATE %s SET class='%s' WHERE `genbank_accession` IN ('%s')" % (table, virus_class, accession_str)
            cursor.execute(sql)
            connection.commit()
            accession = []
            count = 0
    if len(accession) > 0:
        accession_str = "','".join(accession)
        sql = "UPDATE %s SET class='%s' WHERE `genbank_accession` IN ('%s')" % (table, virus_class, accession_str)
        cursor.execute(sql)
        connection.commit()
    infile.close()


if __name__ == '__main__':
    main()
